''' Wiimote.py
	
	Establishes the basic communication mechanism with a Wii remote.
	
	
   Copyright 2008 Corey Tabaka

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
'''

from Hid import *
from State import *
import win32file
import win32event
import pywintypes
import sys
import threading
import time
import traceback
from struct import *
from Queue import Queue

def iif(cond, true, false=0):
	if cond: return true
	else: return false

# VID = Nintendo, PID = Wiimote
VID = 0x057e
PID = 0x0306

REPORT_LENGTH = 22

class ReportMethod(object):
	File				= 0
	Hid					= 1

class InputReport(object):
	Status				= 0x20
	ReadData			= 0x21
	WriteAck			= 0x22
	Buttons				= 0x30
	ButtonsAccel		= 0x31
	IRAccel				= 0x33
	ButtonsExtension	= 0x34
	ExtensionAccel		= 0x35
	IRExtensionAccel	= 0x37

class OutputReport(object):
	None_				= 0x00
	LEDs				= 0x11
	Type				= 0x12
	IR					= 0x13
	Status				= 0x15
	WriteMemory			= 0x16
	ReadMemory			= 0x17
	IR2					= 0x1a

REGISTER_IR						= 0x04b00030
REGISTER_IR_SENSITIVITY_1		= 0x04b00000
REGISTER_IR_SENSITIVITY_2		= 0x04b0001a
REGISTER_IR_MODE				= 0x04b00033

REGISTER_EXTENSION_INIT			= 0x04a40040
REGISTER_EXTENSION_TYPE			= 0x04a400fe
REGISTER_EXTENSION_CALIBRATION	= 0x04a40020

DECODE_TABLE = ''.join([chr( ((x ^ 0x17) + 0x17) & 0xff ) for x in range(0, 256)])

class Wiimote(object):
	ConnectedIds = []
	
	debug_level = 0

	def __init__(self):
		self.id = -1
		self.reportOut = win32file.AllocateReadBuffer(REPORT_LENGTH)
		self.readData = ''
		self.state = WiimoteState()
		self.reportType = InputReport.Status
		self.continuous = False
		self.packets = Queue(0)
		#self.readPackets = Queue(0)
		
		self.reportMethod = ReportMethod.File
		
		self.readEvent = threading.Event()
		self.quitEvent = win32event.CreateEvent(None, 0, 0, None)
		self.quit = False

	def __del__(self):
		#self.Disconnect()
		pass

	def Disconnect(self):
		self.quit = True
		win32event.SetEvent(self.quitEvent)
		try:
			self.handle.Close()
		except:
			pass
		if self.id in Wiimote.ConnectedIds:
			Wiimote.ConnectedIds.remove(self.id)

	def BeginAsyncRead(self):
		pass
	
	def AsyncReadThread(self):
		while not self.quit:
			overlapped = pywintypes.OVERLAPPED()
			overlapped.Offset = 0
			overlapped.hEvent = win32event.CreateEvent(None, 0, 0, None)
			buffer = win32file.AllocateReadBuffer(REPORT_LENGTH)
			buffer[:] = '\x00' * REPORT_LENGTH
			
			rc, data = win32file.ReadFile(self.handle, buffer, overlapped)
			result = win32event.WaitForMultipleObjects((self.quitEvent, overlapped.hEvent), 0, win32event.INFINITE)
			
			if result == win32event.WAIT_OBJECT_0:
				continue
			elif result == win32event.WAIT_OBJECT_0+1:
				self.packets.put(data, True)
			else:
				for key, value in win32event.__dict__.items():
					if value == result:
						result = '%s (%d)' % (key, value)
				
				if Wiimote.debug_level > 0: print 'Received error code: %s' % result
				return
		
		if Wiimote.debug_level > 1: print 'Quitting AsyncReadThread'
	
	def ReadQueue(self):
		while not self.quit:
			try:
				packet = self.packets.get(True, 0.1)
				self.ParseInput(packet)
			except:
				pass
		
		if Wiimote.debug_level > 1: print 'Quitting ReadQueue'
	
	def StartReadThreads(self):
		self.thread0 = threading.Thread(target=self.AsyncReadThread)
		self.thread0.start()
		self.thread1 = threading.Thread(target=self.ReadQueue)
		self.thread1.start()
		self.thread2 = threading.Thread(target=self.ReadQueue)
		self.thread2.start()
	
	def OnReadData(self, data):
		self.ParseInput(data)
		self.BeginAsyncRead()
	
	def ParseInput(self,  data):
		idata = [x for x in unpack('<22B', data)]
		trace = True
		
		type = ord(data[0])
		if type == InputReport.Buttons:
			self.ParseButtons(data)
			trace = False
		elif type == InputReport.ButtonsAccel:
			self.ParseButtons(data)
			self.ParseAccel(data)
			trace = False
		elif type == InputReport.IRAccel:
			self.ParseButtons(data)
			self.ParseAccel(data)
			self.ParseIR(data)
			trace = False
		elif type == InputReport.ButtonsExtension:
			self.ParseButtons(data)
			self.ParseExtension(self.DecodeBuffer(data[3:]))
			trace = False
		elif type == InputReport.ExtensionAccel:
			self.ParseButtons(data)
			self.ParseAccel(data)
			self.ParseExtension(self.DecodeBuffer(data[6:]))
			trace = False
		elif type == InputReport.IRExtensionAccel:
			self.ParseButtons(data)
			self.ParseAccel(data)
			self.ParseIR(data)
			self.ParseExtension(self.DecodeBuffer(data[16:]))
			trace = False
		elif type == InputReport.Status:
			self.ParseButtons(data)
			self.state.Battery = idata[6]
			extension = (idata[3] & 0x02) != 0
			
			#if self.state.Extension != extension:
			self.SetReportType(self.reportType, self.continuous)
			
			self.state.Extension = extension
			
			if Wiimote.debug_level > 0: print 'Battery: %d' % self.state.Battery
			if Wiimote.debug_level > 0: print 'Extension: %s' % self.state.Extension

			if self.state.Extension:
				self.InitializeExtension()
			else:
				self.state.ExtensionType = ExtensionType.None_

			# report extension state changed event
			
			trace = False
		elif type == InputReport.ReadData:
			self.ParseButtons(data)
			self.ParseReadData(data)
			trace = False
		elif type == InputReport.WriteAck:
			pass # format not understood
		else:
			if Wiimote.debug_level > 0: print 'Unknown report type: %02x' % type
		
		if trace:
			if Wiimote.debug_level > 1: print 'Data:   ', ' '.join(['%02x' % x for x in idata])

	def ParseButtons(self, data):
		self.state.ButtonState.update(data)
		if Wiimote.debug_level > 1: print str(self.state.ButtonState)

	def ParseAccel(self, data):
		self.state.AccelState.update(data, self.state.CalibrationInfo)
		if Wiimote.debug_level > 1: print str(self.state.AccelState)

	def ParseIR(self, data):
		self.state.IRState.update(data)
		if Wiimote.debug_level > 1: print str(self.state.IRState)

	def ParseExtension(self, data):
		#print str(self.state.ClassicControllerState)
		if self.state.ExtensionType == ExtensionType.Nunchuk:
			self.state.NunchukState.update(data)
			if Wiimote.debug_level > 0: print str(self.state.NunchukState)
		elif self.state.ExtensionType == ExtensionType.ClassicController:
			self.state.ClassicControllerState.update(data)
			if Wiimote.debug_level > 0: print str(self.state.ClassicControllerState)
	
	def ParseReadData(self, data):
		status = ord(data[3])
		if Wiimote.debug_level > 1: print 'ParseReadData: %08x' % status
		
		if (status & 0x08 ) != 0:
			if Wiimote.debug_level > 0: print 'Error reading from wiimote: Bytes do not exist.'
		elif (status & 0x07) != 0:
			if Wiimote.debug_level > 0: print 'Error reading from wiimote: Attempt to read from write-only registers.'
		else:
			size = status >> 4
			self.readData = data[6:6+size+1]
			self.readEvent.set()
	
	def GetRumbleBit(self):
		if self.state.Rumble:
			return 0x01
		else:
			return 0x00
	
	def ReadData(self, address, size):
		self.readData = '\x00' * size
		self.readEvent.clear()
		
		address |= self.GetRumbleBit() << 24
		report = pack('>BIH', OutputReport.ReadMemory, address, size)
		self.WriteReport(report)
		
		self.readEvent.wait(2.5)
		if self.readEvent.isSet():
			self.readEvent.clear()
			return self.readData
		elif self.reportMethod == ReportMethod.File:
			# this timeout may be an indication that the Bluetooth stack in use
			# doesn't work with the file write method, retry with the HID report
			# method. if the read times out again then there may be a different
			# root cause.
			self.reportMethod = ReportMethod.Hid
			return self.ReadData(address, size)
		else:
			raise Exception('Error reading data from wiimote. Is it connected?')
		
	def WriteData(self, address, data):
		if isinstance(data, int):
			if data <= 0xff:
				data = pack('>B', data)
			elif data <= 0xffff:
				data = pack('>H', data)
			else:
				data = pack('>I', data)
		
		address |= self.GetRumbleBit() << 24
		report = pack('>BIB', OutputReport.WriteMemory, address, len(data))
		report += data
		
		self.WriteReport(report)
		
		time.sleep(0.01)
	
	def ReadCalibration(self):
		data = self.ReadData(0x0016, 7)
		
		if data:
			self.state.CalibrationInfo.update(data)
			if Wiimote.debug_level > 1: print str(self.state.CalibrationInfo)
	
	def DecodeBuffer(self, data):
		return data.translate(DECODE_TABLE)
	
	def InitializeExtension(self):
		if Wiimote.debug_level > 1: print 'InitializeExtension:enter'
		try:
			self.WriteData(REGISTER_EXTENSION_INIT, 0x00)
			
			data = self.ReadData(REGISTER_EXTENSION_TYPE, 2)
			buff = unpack('>BB', data[0:2])
			
			if Wiimote.debug_level > 1: print buff
			
			if buff[0] == ExtensionType.Nunchuk and buff[1] == ExtensionType.Nunchuk:
				self.state.ExtensionType = ExtensionType.Nunchuk
				data = self.DecodeBuffer(self.ReadData(REGISTER_EXTENSION_CALIBRATION, 16))
				self.state.NunchukState.CalibrationInfo.update(data)
				if Wiimote.debug_level > 0: print str(self.state.NunchukState.CalibrationInfo)
			elif buff[0] == ExtensionType.ClassicController and buff[1] == ExtensionType.ClassicController:
				self.state.ExtensionType = ExtensionType.ClassicController
				data = self.DecodeBuffer(self.ReadData(REGISTER_EXTENSION_CALIBRATION, 16))
				self.state.ClassicControllerState.CalibrationInfo.update(data)
				if Wiimote.debug_level > 0: print str(self.state.ClassicControllerState.CalibrationInfo)
			elif buff[0] == 0xff:
				self.state.Extension = False
				self.state.ExtensionType = ExtensionType.None_
		except:
			traceback.print_exc()
		if Wiimote.debug_level > 1: print 'InitializeExtension:exit'

	def WriteReport(self, report):
		report += '\x00' * (22 - len(report))
		if Wiimote.debug_level > 1: print 'Writing:', ' '.join(['%02x' % x for x in unpack('<22B', report)])
		
		if self.reportMethod == ReportMethod.File:
			overlapped = pywintypes.OVERLAPPED()
			overlapped.hEvent = win32event.CreateEvent(None, 0, 0, None)
			self.reportOut[:] = report[:]
			rc, count = win32file.WriteFile(self.handle, self.reportOut, overlapped)
		elif self.reportMethod == ReportMethod.Hid:
			data = (c_char * REPORT_LENGTH)()
			data[:] = report[:]
			HidD_SetOutputReport(self.handle.handle, pointer(data), REPORT_LENGTH)
		else:
			raise Exception('Bogus report method')
		time.sleep(0.01)

	def SetLEDs(self, led1, led2, led3, led4):
		self.state.LEDs.update(led1, led2, led3, led4)
		
		report = pack('<BB', OutputReport.LEDs,
			(iif(led1, 0x10) | iif(led2, 0x20) | iif(led3, 0x40) | iif(led4, 0x80) | self.GetRumbleBit()))
		
		self.WriteReport(report)
	
	def SetRumble(self, on):
		self.state.Rumble = bool(on)
		self.SetLEDs(self.state.LEDs.LED1, self.state.LEDs.LED2, self.state.LEDs.LED3, self.state.LEDs.LED4)
	
	def EnableIR(self, mode):
		self.state.IRState.Mode = mode
		
		report = pack('>BB', OutputReport.IR, 0x04 | self.GetRumbleBit())
		self.WriteReport(report)
		
		report = pack('>BB', OutputReport.IR2, 0x04 | self.GetRumbleBit())
		self.WriteReport(report)
		
		self.WriteData(REGISTER_IR, 0x08)
		
		# high sensitivity
		self.WriteData(REGISTER_IR_SENSITIVITY_1, '\x00\x00\x00\x00\x00\x00\x90\x00\x41')
		self.WriteData(REGISTER_IR_SENSITIVITY_2, '\x40\x00')
		
		self.WriteData(REGISTER_IR_MODE, mode)
	
	def DisableIR(self):
		self.state.IRState.Mode = IRMode.Off
		
		report = pack('>BB', OutputReport.IR, self.GetRumbleBit())
		self.WriteReport(report)
		
		report = pack('>BB', OutputReport.IR2, self.GetRumbleBit())
		self.WriteReport(report)
	
	def SetReportType(self, type, continuous):
		self.reportType = type
		self.continuous = continuous
		
		if type == InputReport.IRAccel:
			self.EnableIR(IRMode.Extended)
		elif type == InputReport.IRExtensionAccel:
			self.EnableIR(IRMode.Basic)
		else:
			self.DisableIR()
		
		report = pack('>BBB', OutputReport.Type, iif(continuous, 0x04) | self.GetRumbleBit(), type)
		self.WriteReport(report)
		
	def GetBatteryLevel(self):
		report = pack('>BB', OutputReport.Status, self.GetRumbleBit())
		self.WriteReport(report)
	
	def GetStatus(self):
		report = pack('>B', OutputReport.Status)
		self.WriteReport(report)

	def Connect(self):
		index = 0
		found = False
		guid = Guid()
		pguid = pointer(guid)

		# get guid of HID class
		HidD_GetHidGuid(pguid)

		# get a handle to all devices that are part of the HID class
		hDevInfo = SetupDiGetClassDevs(pguid, None, 0, DIGCF_DEVICEINTERFACE | DIGCF_PRESENT)

		diData = SP_DEVICE_INTERFACE_DATA()
		diData.cbSize = sizeof(diData)

		while SetupDiEnumDeviceInterfaces(hDevInfo, None, pguid, index, pointer(diData)):
			size = c_uint32()
			SetupDiGetDeviceInterfaceDetail(hDevInfo, pointer(diData), None, 0, pointer(size), None)

			diDetail = SP_DEVICE_INTERFACE_DETAIL_DATA()
			if sizeof(POINTER(c_uint32)) == 8:
				diDetail.cbSize = 8
			else:
				diDetail.cbSize = 5

			if SetupDiGetDeviceInterfaceDetail(hDevInfo, pointer(diData), pointer(diDetail), size, pointer(size), None):
				if Wiimote.debug_level > 0: print index, diDetail.DevicePath

				try:
					handle = win32file.CreateFile(diDetail.DevicePath, win32file.GENERIC_READ | win32file.GENERIC_WRITE,
						win32file.FILE_SHARE_READ | win32file.FILE_SHARE_WRITE, None, win32file.OPEN_EXISTING,
						EFileAttributes.Overlapped, None)
				except:
					traceback.print_exc()
					index += 1
					continue

				attrib = HIDD_ATTRIBUTES()
				attrib.cbSize = sizeof(attrib)

				if HidD_GetAttributes(handle.handle, pointer(attrib)):
					if attrib.VendorID == VID and attrib.ProductID == PID:
						if Wiimote.debug_level > 0: print 'Found a wiimote!'
						if not Wiimote.IsConnected(index):
							if Wiimote.debug_level > 0: print 'Connecting to wiimote %d' % index
							self.id = index
							Wiimote.ConnectedIds.append(index)
							found = True
							self.handle = handle
							
							self.StartReadThreads()

							self.BeginAsyncRead()
							#self.GetStatus()
							self.ReadCalibration()

							break
						else:
							if Wiimote.debug_level > 0: print 'Already connected to wiimote %d' % index
							handle.Close()
					else:
						handle.Close()

			index += 1

		SetupDiDestroyDeviceInfoList(hDevInfo)

		if not found:
			raise Exception('No wiimote found')

	@staticmethod
	def IsConnected(id):
		return id in Wiimote.ConnectedIds

if __name__=="__main__":
	import sys

	try:
		Wiimote.debug_level = 2
		wiimote = Wiimote()
		wiimote.Connect()
		wiimote.GetStatus()
		#wiimote.DisableIR()
		wiimote.SetReportType(InputReport.Buttons, False)
		#wiimote.SetLEDs(1, 0, 1, 0)
		#time.sleep(0.5)
		#wiimote.SetLEDs(0, 1, 0, 1)
		#time.sleep(0.5)
		#wiimote.SetLEDs(1, 1, 0, 0)
		#time.sleep(0.5)
		#wiimote.SetLEDs(0, 0, 1, 1)
		#time.sleep(0.5)
		wiimote.SetLEDs(1, 0, 0, 0)
		wiimote.SetReportType(InputReport.IRExtensionAccel, False)
		#wiimote.SetReportType(InputReport.ButtonsExtension, False)
		#wiimote.SetReportType(InputReport.ButtonsAccel, False)
	except:
		wiimote.Disconnect()
		raise

	raw_input('Press enter to quit\n')
	
	wiimote.DisableIR()
	wiimote.SetReportType(InputReport.Buttons, False)
	wiimote.Disconnect()
